"""
@Time: 2021/1/12 下午 9:16
@Author: jinzhuan
@File: cluener.py
@Desc: 
"""
import os
from ..loader import Loader
import json


class ClueNERLoader(Loader):

    def __init__(self):
        super().__init__()
        self.label_set.add('O')

    def _load(self, path):
        dataset = []
        with open(path) as f:
            for line in f:
                data = json.loads(line)
                label = data['label']
                for key, value in label.items():
                    self.label_set.add('B-{}'.format(key))
                    self.label_set.add('I-{}'.format(key))
                dataset.append(data)
        return dataset

    def load_all(self, path):
        train_path = os.path.join(path, 'train.json')
        dev_path = os.path.join(path, 'dev.json')
        return self._load(train_path), self._load(dev_path)
